import sys
sys.path.append('C:/Users/Hu/Dropbox/Research/PythonWork/Cancer/src/STAT/')

import numpy as np
import csv
from datetime import datetime
from time import strftime
import matplotlib.pyplot as plt
import math
import random
from scipy import stats
import ols

'''
Global and Local Empirical Bayes Smoothers with Gamma Model
'''

def getCurTime():
    """
    get current time
    Return value of the date string format(%Y-%m-%d %H:%M:%S)
    """
    format='%Y-%m-%d %H:%M:%S'
    sdate = None
    cdate = datetime.now()
    try:
        sdate = cdate.strftime(format)
    except:
        raise ValueError
    return sdate

def build_data_list(inputCSV):
    sKey = []
    fn = inputCSV
    f = open(inputCSV)
    #ra = csv.DictReader(file(fn), dialect="excel")
    ra = csv.DictReader(f, dialect="excel")
    
    for record in ra:
        #print record[ra.fieldnames[0]], type(record[ra.fieldnames[-1]])
        for item in ra.fieldnames:
            temp = int(float(record[item]))
            sKey.append(temp)
    sKey = np.array(sKey)
    sKey.shape=(-1,len(ra.fieldnames))
    return sKey

def geneBpChoices(tmin, tmax, num):
    tmin = float(tmin)
    tmax = float(tmax)
    num = int(num)
    step = (tmax - tmin)/num
    tbpchoices = []
    for i in range(num):
        tbpchoices.append(tmin+i*step)
    return tbpchoices

def pickBreakpoint(response, x1, predictor):
    #print int(min(predictor))*10, int(max(predictor)+1)*10, int(max(predictor) - min(predictor) + 1)/2
    #bpChoices = geneBpChoices(min(predictor), max(predictor), 20)
    results = np.zeros((len(bpChoices)-1, 2))
    print bpChoices
    predictor = np.array(predictor)
    
    for i in range(len(bpChoices)-1):
        #print i
        #print type((predictor - bpChoices[i]))
        x2star = (predictor - bpChoices[i]) * np.greater(predictor, bpChoices[i])
        tempPredictor = np.array(zip(x1, predictor, x2star))
        #fileLoc = filePath + 'temp.csv'
        #np.savetxt(fileLoc, tempPredictor, delimiter=',', fmt = '%s')
        tempmodel = ols.ols(response, tempPredictor,'y',['F1F2', 'dist', 'diststar'])
        results[i,0] = i
        #results[i,1] = tempmodel.sse
        results[i,1] = tempmodel.R2
    print results

    optBP = int(results[np.argmax(results, axis = 0)[1],0])
    print 'Optimal Index:', optBP
    print 'Optimal changepoint: ', bpChoices[optBP], ' exp value: ', np.exp(bpChoices[optBP]), ' with R2 = ', results[optBP, 1]

    #x2star = (predictor - bpChoices[optBP]) * np.greater(predictor, bpChoices[optBP])
    #optPredictor = np.array(zip(predictor, x2star))
    #optmodel = ols.ols(response, optPredictor,'y',['x1', 'x2'])
    x2star = (predictor - bpChoices[optBP]) * np.greater(predictor, bpChoices[optBP])
    optPredictor = np.array(zip(x1, predictor, x2star))
    optmodel = ols.ols(response, optPredictor,'y',['F1F2', 'dist', 'diststar'])
    #print optmodel.b
    #return bpChoices[optBP], results, optmodel, optmodel.b[0]+optmodel.b[1]*predictor+optmodel.b[2]*x2star
    y_hat = np.dot(optmodel.x, optmodel.b)
    fileLoc = filepath + 'pieceR_model2_y_hat.csv'
    #np.savetxt(fileLoc, y_hat, delimiter=',', fmt = '%s')
    
    print optmodel.summary()
    print 'MSE =', optmodel.sse
    print 'before bp, b0 =',optmodel.b[0], ', b1 =', optmodel.b[1], ', bd =', optmodel.b[2]
    print 'after bp, b0 =', optmodel.b[0] - optmodel.b[3] * bpChoices[optBP], ', b1 =', optmodel.b[1], ', bd =', optmodel.b[2] + optmodel.b[3]

    calpredictedvalue(zip(x1, predictor), bpChoices[optBP], zip(optmodel.b, [optmodel.b[0] - optmodel.b[3] * bpChoices[optBP], optmodel.b[1], optmodel.b[2] + optmodel.b[3]]), 'exp_inoutflow_model2A.csv')
    return results, y_hat

def calexpectation(b0, b1, bd):
    f12 = 455401.0 * 198508
    d = 2219333.0
    print math.pow(math.e, b0) * math.pow(f12, b1) * math.pow(d, bd)

def pickBreakpointV2(response, x1, predictor):
    #print int(min(predictor))*10, int(max(predictor)+1)*10, int(max(predictor) - min(predictor) + 1)/2
    #bpChoices = geneBpChoices(min(predictor), max(predictor), 20)
    results = np.zeros((len(bpChoices)-1, 2))
    print bpChoices
    
    for i in range(len(bpChoices)-1):
        print i
        x2star = (predictor - bpChoices[i]) * np.greater(predictor, bpChoices[i])
        x1star = x1 * np.greater(predictor, bpChoices[i]) 
        tempPredictor = np.array(zip(x1, x1star, predictor, x2star))
        #fileLoc = filePath + 'temp.csv'
        #np.savetxt(fileLoc, tempPredictor, delimiter=',', fmt = '%s')
        #print tempPredictor
        tempmodel = ols.ols(response, tempPredictor,'y',['F1F2', 'F1F2star', 'dist', 'diststar'])
        results[i,0] = i
        #results[i,1] = tempmodel.sse
        results[i,1] = tempmodel.R2

    optBP = int(results[np.argmax(results, axis = 0)[1],0])
    print 'Optimal Index:', optBP
    print 'Optimal changepoint: ', bpChoices[optBP], ' exp value: ', np.exp(bpChoices[optBP]), ' with R2 = ', results[optBP, 1]

    #x2star = (predictor - bpChoices[optBP]) * np.greater(predictor, bpChoices[optBP])
    #optPredictor = np.array(zip(predictor, x2star))
    #optmodel = ols.ols(response, optPredictor,'y',['x1', 'x2'])
    x1star = x1 * np.greater(predictor, bpChoices[optBP])
    x2star = (predictor - bpChoices[optBP]) * np.greater(predictor, bpChoices[optBP])
    optPredictor = np.array(zip(x1, x1star, predictor, x2star))
    optmodel = ols.ols(response, optPredictor,'y',['F1F2', 'F1F2star', 'dist', 'diststar'])
    
    #return bpChoices[optBP], results, optmodel, optmodel.b[0]+optmodel.b[1]*predictor+optmodel.b[2]*x2star
    print results, optmodel.b
    print optmodel.summary()
    return results

def separateData(data1, data2, dataC, threshold):
    if (len(data1) <> len(data2)) or (len(data1) <> len(dataC)):
        print 'Error in input data length! --- seperateData.error.1'

    data1Left = []
    data1Right = []
    data2Left = []
    data2Right = []
    dataleftIdx = []
    datarightIdx = []
    
    for i in range(len(data1)):
        if dataC[i] < threshold:
            data1Left.append(data1[i])
            data2Left.append(data2[i])
            dataleftIdx.append(i)
        else:
            data1Right.append(data1[i])
            data2Right.append(data2[i])
            datarightIdx.append(i)

    return np.array(data1Left), np.array(data1Right), np.array(data2Left), np.array(data2Right), dataleftIdx, datarightIdx

def calculateR2(tleftmodel, trightmodel, totalmean):
    #self.sse = dot(self.e,self.e)/self.df_e
    #tsse = np.dot(tmodel.e, tmodel.e)/tmodel.df_e
    #tsse = np.dot(tmodel.e, tmodel.e)
    tsse = tleftmodel.sse * tleftmodel.df_e + trightmodel.sse * trightmodel.df_e
    
    #y_hat = np.dot(tmodel.x,tmodel.b)
    tsst = np.dot(tleftmodel.y - totalmean, tleftmodel.y - totalmean)
    tsst = tsst + np.dot(trightmodel.y - totalmean, trightmodel.y - totalmean)
    
    #print tleftmodel.sse * tleftmodel.df_e, trightmodel.sse * trightmodel.df_e, np.dot(tleftmodel.y - totalmean, tleftmodel.y - totalmean), np.dot(trightmodel.y - totalmean, trightmodel.y - totalmean)
    #print 1 - (tleftmodel.sse * tleftmodel.df_e)/np.dot(tleftmodel.y - totalmean, tleftmodel.y - totalmean)
    #print 1 - (trightmodel.sse * trightmodel.df_e)/np.dot(trightmodel.y - totalmean, trightmodel.y - totalmean)
    
    return 1-tsse/tsst
    #print ((tleftmodel.nobs+trightmodel.nobs-1.0), (tleftmodel.nobs+trightmodel.nobs-tleftmodel.ncoef-trightmodel.ncoef)), ((tleftmodel.nobs+trightmodel.nobs-1.0)/(tleftmodel.nobs+trightmodel.nobs-tleftmodel.ncoef-trightmodel.ncoef))
    #return 1 - (tsse/tsst)*((tleftmodel.nobs+trightmodel.nobs-1.0)/(tleftmodel.nobs+trightmodel.nobs-tleftmodel.ncoef-trightmodel.ncoef))

def simpleRegression(response, predictor, predictorName):
    model = ols.ols(response, predictor,'y',predictorName)
    #model = ols.ols(response, predictor)
    print model.summary()
    print model.b
    print predictor
    #print 'b0 =',model.b[0], ', b1 =', model.b[1], ', bd =', model.b[2], ', MSE =', model.sse
    return np.dot(model.x, model.b)

def separateRegression(response, predictor, sepData, bpChoices, predictorName):
    results = np.zeros((len(bpChoices)-1, 4))

    for bpid in range(len(bpChoices)-1):
        print bpid
        responseLeft, responseRight, predictorLeft, predictorRight, dataleftIdx, datarightIdx = separateData(response, predictor, sepData, bpChoices[bpid])
        #print np.mean(responseLeft), np.mean(responseRight)
        #print predictorLeft, predictorRight
        leftmodel = ols.ols(responseLeft, predictorLeft,'y',predictorName)
        rightmodel = ols.ols(responseRight, predictorRight,'y',predictorName)
        results[bpid,0] = bpid
        results[bpid,1] = leftmodel.R2adj
        results[bpid,2] = rightmodel.R2adj
        #results[bpid,3] = 1 - (leftmodel.e.var() + rightmodel.e.var())/(leftmodel.y.var() + rightmodel.y.var())
        results[bpid,3] = calculateR2(leftmodel, rightmodel, np.mean(response))
        #results[bpid,3] = (leftmodel.R2 + rightmodel.R2)/2

    print results
        
    optBP = int(results[np.argmax(results, axis = 0)[-1],0])

    responseLeft, responseRight, predictorLeft, predictorRight, dataleftIdx, datarightIdx = separateData(response, predictor, sepData, bpChoices[optBP])
    leftmodel = ols.ols(responseLeft, predictorLeft,'y',predictorName)
    rightmodel = ols.ols(responseRight, predictorRight,'y',predictorName)
    #print calculateR2(leftmodel)
    #optmodel = ols.ols(response, optPredictor,'y',['F1F2', 'F1F2star', 'dist', 'diststar'])
    yhatL = np.dot(leftmodel.x, leftmodel.b)
    yhatR = np.dot(rightmodel.x, rightmodel.b)
    yhat = np.zeros(len(response))

    for i in range(len(yhatL)):
        yhat[dataleftIdx[i]] = yhatL[i]

    for i in range(len(yhatR)):
        yhat[datarightIdx[i]] = yhatR[i]

    yhat = np.exp(yhat)
    fileLoc = filepath + 'separateR_model2_y_hat.csv'
    #np.savetxt(fileLoc, yhat, delimiter=',', fmt = '%s')
    print 'Optimal Index:', optBP
    print 'Optimal changepoint: ', bpChoices[optBP], ' exp value: ', np.exp(bpChoices[optBP]), ' with R2 = ', results[optBP, -1]
   
    #return bpChoices[optBP], results, optmodel, optmodel.b[0]+optmodel.b[1]*predictor+optmodel.b[2]*x2star
    #print results, optmodel.b
    print '----------------------------- left model -----------------------------'
    print leftmodel.summary()
    print '----------------------------- right model -----------------------------'
    print rightmodel.summary()

    print 'Optimal Index:', optBP
    print 'Optimal changepoint: ', bpChoices[optBP], ' exp value: ', np.exp(bpChoices[optBP]), ' with R2 = ', results[optBP, -1]
   

    print 'before bp, b0 =',leftmodel.b[0], ', b1 =', leftmodel.b[1], ', bd =', leftmodel.b[2], ', MSE =', leftmodel.sse
    print 'after bp, b0 =', rightmodel.b[0], ', b1 =', rightmodel.b[1], ', bd =', rightmodel.b[2], ', MSE =', rightmodel.sse

    #calpredictedvalue(predictor, bpChoices[optBP], zip(leftmodel.b, rightmodel.b), 'exp_inoutflow_model2B.csv')
    #calconfidenceinterval(predictorLeft, predictorRight, [leftmodel.sse, rightmodel.sse], response, predictor, bpChoices[optBP], zip(leftmodel.b, rightmodel.b), 'ci_model2B.csv')
    return results, yhat

def calconfidenceinterval(leftx, rightx, mse, y, x, bp, b,filename):
    leftx = np.array(zip(np.ones(len(leftx)), leftx[:,0], leftx[:,1]))
    rightx = np.array(zip(np.ones(len(rightx)), rightx[:,0], rightx[:,1]))

    templeftx = np.linalg.inv(np.dot(leftx.T, leftx))
    temprightx = np.linalg.inv(np.dot(rightx.T, rightx))

    b = np.array(b)
    t = [stats.t.ppf(0.975, len(leftx) - len(leftx[0,:])), stats.t.ppf(0.975, len(rightx) - len(rightx[0,:]))]

    yhat = []
    i = 0
    #print len(x)
    for item in x:
        if item[1] < bp:
            bpflag = 0
            tempx = templeftx
        else:
            bpflag = 1
            tempx = temprightx
        temp = b[0, bpflag] + b[1, bpflag] * item[0] + b[2, bpflag] * item[1]
        xh = np.array([1, item[0], item[1]])
        #print np.dot(np.dot(xh, tempx), xh.T), mse[bpflag]*(np.dot(np.dot(xh, tempx), xh.T)
        se = math.pow(mse[bpflag]*(np.dot(np.dot(xh, tempx), xh.T)), 0.5) * t[bpflag]
        pise = math.pow(mse[bpflag]*(1 + np.dot(np.dot(xh, tempx), xh.T)), 0.5) * t[bpflag]
        cil = 0
        cih = 0
        pil = 0
        pih = 0
        if y[i] < (temp-se):
            cil = 1
        if y[i] > (temp+se):
            cih = 1
        if y[i] < (temp-pise):
            pil = 1
        if y[i] > (temp+pise):
            pih = 1
        yhat.append([temp, temp-se, temp+se, temp-pise, temp+pise,cil,cih,pil,pih])
        i += 1

    yhat = np.exp(yhat)
    fileLoc = filepath + filename   #'exp_inoutflow_model2A.csv'
    headerstr = 'yhat,cileft,ciright,pileft,piright,cil,cih,pil,pih'
    #print fileLoc
    np.savetxt(fileLoc, yhat, delimiter=',', header = headerstr, fmt = '%f')
    
    #plt.plot(range(len(y)), y, 'g', label='Observed')
    #plt.plot(range(len(y)), yhat[:,0], 'b', label='Predicted')
    #plt.plot(range(len(y)), yhat[:,1], 'b', label='CI Left')
    #plt.plot(range(len(y)), yhat[:,2], 'b', label='CI Right')
    #plt.plot(range(len(y)), yhat[:,3], 'r', label='PI Left')
    #plt.plot(range(len(y)), yhat[:,4], 'r', label='PI Right')
    #plt.show()
    

def calpredictedvalue(x, bp, b, filename):
    # x: [f1f2, dist]
    # b = [[b00, b10], [b01, b11], [b02, b12]]
    print b
    yhat = []
    b = np.array(b)
    bpflag = 1  # 0: before, 1: after
    for item in x:
        if item[1] < bp:
            bpflag = 0
        else:
            bpflag = 1
        temp = b[0, bpflag] + b[1, bpflag] * item[0] + b[2, bpflag] * item[1]
        #temp = np.dot(item, b[:,bpflag])
        yhat.append(temp)

    yhat = np.array(yhat)
    yhat = np.exp(yhat)
    #print yhat
    #fileLoc = filepath + 'model2b_yhat.csv'
    #np.savetxt(fileLoc, yhat, delimiter=',', fmt = '%f')
    expinoutflow = np.zeros((len(inoutflow), 2)) # [inflow, outflow, obs-exp inflow, obs-exp outflow]
    #print len(yhat), len(data)
    #print zip(data, yhat)
    i = 0
    for item in data:
        expinoutflow[int(item[1]), 1] += yhat[i]
        expinoutflow[int(item[2]), 0] += yhat[i]
        i += 1

    fileLoc = filepath + filename   #'exp_inoutflow_model2A.csv'
    headerstr = 'inflow,outflow'
    #np.savetxt(fileLoc, expinoutflow, delimiter=',', header = headerstr, fmt = '%f')
    
    
    
#--------------------------------------------------------------------------
#MAIN

if __name__ == "__main__":
    print '===================================================='
    print "begin at " + getCurTime()

    #calexpectation(-1.5529, 1.1522, -1.4576)
    #calexpectation(-14.0855, 1.1522, -0.5477)
    #a.append([5])
    #print geneBpChoices(1.4, 7.9, 20)

    #filePath = 'C:/_DATA/migration_census_2000/Aggregation/'
    #file = filePath + 'census_county_migration_aggregation_dist.csv' #[dist, orgid, desid, flowvol]
    #filepath = 'C:/_DATA/migration_census_2000/REDCAPLUS_100Region/'
    #file = filepath + 'census_migration_100regs_dis.csv'
    filepath = 'C:/_DATA/migration_census_2000/STATE/'
    file = filepath + 'census_migration_state_dist.csv'
    data = build_data_list(file)    #[dist, orgid, destid, flowvol]

    #inoutflowfile = filePath + 'census_county_migration_aggregation_inoutflow.csv'    #[inflow, outflow]
    #inoutflowfile = filepath + 'census_migration_100regs_inoutflow.csv'
    inoutflowfile = filepath + 'census_migration_state_pop2000.csv'
    inoutflow = build_data_list(inoutflowfile)

    #West FID:
    westid = [1,2,3,4,10,13,14,16,21,23,24,25,26,29,32,34,35,39,41,42,45,48]

    #East FID:
    eastid = [0,5,6,7,8,9,11,12,15,17,18,19,20,22,27,28,30,31,33,36,37,38,40,43,44,46,47]

    dist = data[:,0]

    usedistgraph = 1
    if(usedistgraph):
        distgraphcsv = filepath + 'graphdistance.csv'
        distgraph = build_data_list(distgraphcsv)
        dist = []
        for item in data:
            dist.append(distgraph[int(item[1]), int(item[2])])

    lndist = np.log(dist)
    
    x1 = []
    for item in data:
        x1.append(np.log(float(inoutflow[item[1],1])*inoutflow[item[2],0]))

    bpChoices = np.log(range(min(dist)+2, max(dist)+1))
    #bpChoices = range(600000,3100000,100000)
    #bpChoices = np.log(range(800000,2500000,100000))
    print bpChoices
    #print np.log(bpChoices)

    #inputPredictor = np.array(zip(x1, np.log(dist)))

    lenid = len(np.unique(data[:,1]))
    #indicators = np.zeros((len(data), lenid*2))

    i = 0
    output = np.array([])

    '''
    for item in data:
        output = np.append(output, [item[-1], x1[i], dist[i]])
        temp = np.zeros(lenid*2)
        temp[int(item[1])*2] = 1
        temp[int(item[2])*2+1] = 1
        output = np.append(output, temp)
        i += 1
    '''
    '''
    for item in data:
        output = np.append(output, [math.log(item[-1]), x1[i], dist[i]])
        temp = np.zeros(lenid)
        temp[int(item[1])] = 1
        temp[int(item[2])] = 1
        output = np.append(output, temp)
        temp = np.zeros(lenid)
        temp[int(item[1])] = inoutflow[item[1],1]
        temp[int(item[2])] = inoutflow[item[2],0]
        output = np.append(output, temp)
        temp = [0,0] #oe,de
        if int(item[1]) in eastid:
            temp[0] = 1
        #if int(item[1]) in westid:
            #temp[2] = 1
        if int(item[2]) in eastid:
            temp[1] = 1
        #if int(item[2]) in westid:
            #temp[3] = 1
        #if (temp[0] == 1) and (temp[2] == 1):
            #print "Error!!!"
        #if (temp[1] == 1) and (temp[3] == 1):
            #print "Error!!!"
        output = np.append(output, temp)
        i += 1
    '''
    lninoutflow = np.log(inoutflow)
    for item in data:
        output = np.append(output, [math.log(item[-1]), lninoutflow[item[1],1], lninoutflow[item[2],0], lninoutflow[item[1],1] * lninoutflow[item[1],1], lninoutflow[item[2],0] * lninoutflow[item[2],0], lndist[i]])
        i += 1
    output.shape = (len(data), -1)
    #output: [logflowvol, logp1p2, logdist, od0, od1, ..., od48, pop0, pop1, ... pop48, oe, de]


    #plt.scatter(dist, data[:,-1], marker = 'o')
    #plt.scatter(np.log(dist), np.log(data[:,-1]), marker = 'o')
    #plt.loglog(dist, data[:,-1], marker = 'o', linestyle = '')
    '''
    totalflow = np.sum(data[:,-1])
    expdata = np.exp(x1)/totalflow
    avgExp = np.average(expdata)
    ajustOb = []
    for ob, exp in zip(data[:,-1], expdata):
        ajustOb.append(ob/(exp/avgExp))
    '''
    '''
    xx = dist
    yy = data[:,-1]
    #yy = ajustOb
    plt.subplot(221)
    plt.scatter(xx, yy)
    #plt.plot(xx, yy)
    plt.xlabel('graph distance')
    plt.ylabel('flow volume')
    plt.title('Original Data')
    plt.subplot(222)
    plt.scatter(xx, np.log(yy))
    #plt.plot(xx, np.log(yy))
    plt.xlabel('graph distance')
    plt.ylabel('log(flow volume)')
    plt.title('log(y)')
    plt.subplot(223)
    plt.scatter(np.log(xx), np.log(yy))
    #plt.plot(np.log(xx), np.log(yy))
    plt.xlabel('log(graph distance)')
    plt.ylabel('log(flow volume)')
    plt.title('log(x), log(y)')
    plt.subplot(224)
    plt.loglog(xx, yy, marker = 'o', linestyle = '')
    plt.xlabel('graph distance')
    plt.ylabel('flow volume')
    plt.title('loglog plot')
    plt.xlim(xmin = -0.5, xmax = 14)
    plt.show()
    '''

    #inputPredictor = output[:,[1,2,-2,-1]]
    inputPredictor = output[:,[1,2,3,4,5]]
    #inputPredictorName = ['logp1p2', 'logdist','oe','de']
    inputPredictorName = ['lnp1', 'lnp2', 'lnp12', 'lnp22', 'lndist']
    #inputPredictorName = ['lnp1', 'lnp2', 'lnp12', 'lndist']
    #yhat = simpleRegression(output[:,0], inputPredictor, inputPredictorName)
    result, yhat = separateRegression(output[:,0], inputPredictor, lndist, bpChoices, inputPredictorName)
    #result = pickBreakpoint(np.log(data[:,-1]), x1, dist)

    #fileLoc = filepath + 'piecewiseR_model2_results.csv'
    #fileLoc = filepath + 'separateR_model2_results.csv'
    headerstr = 'exp'
    fileLoc = filepath + 'temp.csv'
    #np.savetxt(fileLoc, yhat, delimiter=',', header = headerstr, fmt = '%s')
    